import os
import sys
import asyncio
from multiprocessing import Process, Manager
from platform import system

import websockets
from config import ServerConfig as Config
from util.server_cosmic import Cosmic, console
from util.server_check_model import check_model
from util.server_ws_recv import ws_recv
from util.server_ws_send import ws_send
from util.server_init_recognizer import init_recognizer
from util.empty_working_set import empty_current_working_set

# Import logging system
from util.logging import setup_logging, get_logger, SystemHealthLogger, PerformanceTracker

BASE_DIR = os.path.dirname(__file__); os.chdir(BASE_DIR)    # 确保 os.getcwd() 位置正确，用相对路径加载模型

async def main():
    # Initialize logging system
    setup_logging(profile=os.environ.get('CAPSWRITER_LOG_PROFILE', 'production'))
    logger = get_logger('capswriter.server')
    
    # Initialize system health monitoring
    health_logger = SystemHealthLogger()
    health_logger.log_startup_info()
    health_logger.start_monitoring(interval=300)  # Monitor every 5 minutes
    
    # Initialize performance tracking
    perf_tracker = PerformanceTracker('capswriter.server.performance')
    perf_tracker.start_operation('server_startup')

    logger.info("Starting CapsWriter Offline Server")
    logger.info(f"Base directory: {BASE_DIR}")
    logger.info(f"Server address: {Config.addr}:{Config.port}")

    # 检查模型文件
    try:
        check_model()
        logger.info("Model files validated successfully")
    except Exception as e:
        logger.error("Model validation failed", exc_info=e)
        raise

    console.line(2)
    console.rule('[bold #d55252]CapsWriter Offline Server'); console.line()
    console.print(f'项目地址：[cyan underline]https://github.com/HaujetZhao/CapsWriter-Offline', end='\n\n')
    console.print(f'当前基文件夹：[cyan underline]{BASE_DIR}', end='\n\n')
    console.print(f'绑定的服务地址：[cyan underline]{Config.addr}:{Config.port}', end='\n\n')

    # 跨进程列表，用于保存 socket 的 id，用于让识别进程查看连接是否中断
    Cosmic.sockets_id = Manager().list()

    # 负责识别的子进程
    try:
        logger.info("Starting recognition process")
        recognize_process = Process(target=init_recognizer,
                                    args=(Cosmic.queue_in,
                                          Cosmic.queue_out,
                                          Cosmic.sockets_id),
                                    daemon=True)
        recognize_process.start()
        logger.info(f"Recognition process started with PID: {recognize_process.pid}")
        
        # Wait for recognition process to be ready
        Cosmic.queue_out.get()
        logger.info("Recognition process initialized successfully")
        
    except Exception as e:
        logger.error("Failed to start recognition process", exc_info=e)
        raise

    console.rule('[green3]开始服务')
    console.line()

    # Log server startup completion
    startup_time = perf_tracker.end_operation('server_startup')
    logger.info(f"Server startup completed in {startup_time:.2f} seconds")

    # 清空物理内存工作集
    if system() == 'Windows':
        try:
            empty_current_working_set()
            logger.debug("Memory working set cleared")
        except Exception as e:
            logger.warning("Failed to clear memory working set", exc_info=e)

    def select_subprotocol(connection, subprotocols):
        """Custom subprotocol selector that allows connections with or without subprotocol"""
        if "binary" in subprotocols:
            return "binary"
        # Allow connection even without subprotocol
        return None
    
    try:
        # 负责接收客户端数据的 coroutine
        logger.info("Starting WebSocket server")
        recv = websockets.serve(ws_recv,
                                Config.addr,
                                Config.port,
                                subprotocols=["binary"],
                                select_subprotocol=select_subprotocol,
                                max_size=None)

        # 负责发送结果的 coroutine
        send = ws_send()
        
        logger.info("Server ready to accept connections")
        await asyncio.gather(recv, send)
        
    except Exception as e:
        logger.error("Server operation failed", exc_info=e)
        raise
    finally:
        # Stop health monitoring
        health_logger.stop_monitoring()
        logger.info("Server shutdown completed")


def init():
    logger = None
    try:
        asyncio.run(main())
    except KeyboardInterrupt:           # Ctrl-C 停止
        if logger:
            logger.info("Server shutdown requested by user (Ctrl-C)")
        console.print('\n再见！')
    except OSError as e:                # 端口占用
        if logger:
            logger.error(f"Network error: {e}")
        else:
            # Fallback logging if main logger not initialized
            setup_logging()
            logger = get_logger('capswriter.server')
            logger.error(f"Network error during startup: {e}")
        
        if "Address already in use" in str(e) or "被占用" in str(e):
            console.print(f'端口被占用，可能已有服务端在运行：{e}', style='bright_red')
            console.print('请检查是否已有服务端实例运行，或更改配置中的端口号', style='yellow')
            if logger:
                logger.error(f"Port {Config.port} already in use, another server instance may be running")
        else:
            console.print(f'网络错误：{e}', style='bright_red')
    except Exception as e:
        if logger:
            logger.critical("Unexpected server error", exc_info=e)
        else:
            # Fallback logging
            setup_logging()  
            logger = get_logger('capswriter.server')
            logger.critical(f"Critical error during startup: {e}", exc_info=e)
        
        console.print(f'未知错误：{e}', style='bright_red')
    finally:
        try:
            if hasattr(Cosmic, 'queue_out') and Cosmic.queue_out:
                Cosmic.queue_out.put(None)
            if logger:
                logger.info("Server cleanup completed")
        except Exception as cleanup_error:
            if logger:
                logger.error("Error during cleanup", exc_info=cleanup_error)
        sys.exit(0)
     
        
if __name__ == "__main__":
    init()
